package com.hk.core.elasticsearch.repository;

import co.elastic.clients.elasticsearch.core.search.FieldSuggester;
import co.elastic.clients.elasticsearch.core.search.Suggester;
import com.hk.commons.query.Order;
import com.hk.commons.query.QueryModel;
import com.hk.commons.query.SortUtils;
import com.hk.commons.query.page.QueryPage;
import com.hk.commons.query.page.SimpleQueryPage;
import com.hk.commons.util.BeanUtils;
import com.hk.commons.util.CollectionUtils;
import com.hk.core.elasticsearch.query.Condition;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Persistable;
import org.springframework.data.elasticsearch.client.elc.NativeQuery;
import org.springframework.data.elasticsearch.client.elc.NativeQueryBuilder;
import org.springframework.data.elasticsearch.client.elc.Queries;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.document.Document;
import org.springframework.data.elasticsearch.core.query.Criteria;
import org.springframework.data.elasticsearch.core.query.CriteriaQuery;
import org.springframework.data.elasticsearch.core.query.Query;
import org.springframework.data.elasticsearch.core.query.UpdateQuery;
import org.springframework.data.elasticsearch.core.suggest.response.Suggest;
import org.springframework.data.elasticsearch.repository.support.ElasticsearchEntityInformation;
import org.springframework.data.elasticsearch.repository.support.SimpleElasticsearchRepository;

import java.io.Serializable;
import java.util.*;

/**
 * 自定义扩展 Repository, 添加自定义查询方法
 *
 * @author Kevin
 * @date 2019/3/11 9:00
 */
public class SimpleBaseElasticsearchRepository<T extends Persistable<ID>, ID extends Serializable>
        extends SimpleElasticsearchRepository<T, ID>
        implements BaseElasticsearchRepository<T, ID> {

    public SimpleBaseElasticsearchRepository(ElasticsearchEntityInformation<T, ID> metadata,
                                             ElasticsearchOperations elasticsearchOperations) {
        super(metadata, elasticsearchOperations);
    }

    @Override
    public QueryPage<T> queryForPage(QueryModel<T> queryModel) {
        var map = BeanUtils.beanToMap(queryModel.getParam());
        var criteria = new Criteria();
        if (CollectionUtils.isNotEmpty(map)) {
            map.forEach((key, value) -> criteria.and(Criteria.where(key).is(value)));
        }
        var query = new CriteriaQuery(criteria, PageRequest.of(queryModel.getStartRowIndex(),
                queryModel.getPageSize(), SortUtils.toSort(queryModel.getOrders())));
        var hits = operations.search(query, entityClass);
        return new SimpleQueryPage<>(queryModel, hits.get()
                .map(SearchHit::getContent)
                .toList(), hits.getTotalHits());
    }

    @Override
    public QueryPage<T> queryForPage(List<Condition> conditions, int pageIndex, int pageSize, Order... orders) {
        var query = new CriteriaQuery(new Criteria(),
                PageRequest.of(pageIndex, pageSize, SortUtils.toSort(orders)));
        Condition.addCriteria(query, conditions);
        var hits = operations.search(query, entityClass);
        return new SimpleQueryPage<>(hits.get()
                .map(SearchHit::getContent)
                .toList(), hits.getTotalHits(), pageIndex, pageSize);
    }

    @Override
    public long count(T t) {
        var criteria = new Criteria();
        var map = BeanUtils.beanToMap(t);
        if (CollectionUtils.isNotEmpty(map)) {
            map.forEach((key, value) -> criteria.and(Criteria.where(key).is(value)));
        }
        return operations.count(new CriteriaQuery(criteria), entityClass);
    }

    @Override
    public void deleteByIds(Iterable<ID> ids) {
        if (CollectionUtils.isNotEmpty(ids)) {
            List<String> list = new ArrayList<>();
            ids.forEach(item -> list.add(String.valueOf(item)));
            var idsQuery = Queries.idsQuery(list);
            operations.delete(idsQuery, entityInformation.getIndexCoordinates());
        }
    }

    @Override
    public List<T> findAll(List<Condition> conditions, Order... orders) {
        var query = new CriteriaQuery(new Criteria());
        Condition.addCriteria(query, conditions);
        query.addSort(SortUtils.toSort(orders));
        return operations.search(query, entityClass)
                .get()
                .map(SearchHit::getContent)
                .toList();
    }

    @Override
    public void bulkUpdate(Collection<T> entities) {
        if (CollectionUtils.isNotEmpty(entities)) {
            var list = new ArrayList<UpdateQuery>(entities.size());
            entities.forEach(item -> {
                var id = item.getId();
                if (Objects.nonNull(id)) {
                    var updateMap = BeanUtils.beanToMap(item, "class", "new", "id");
                    if (CollectionUtils.isNotEmpty(updateMap)) {
                        list.add((UpdateQuery.builder(String.valueOf(id))).withDocument(Document.from(updateMap)).build());
                    }
                }
            });
            operations.bulkUpdate(list, entityInformation.getIndexCoordinates());
        }
    }

    @Override
    public List<String> suggest(FieldSuggester suggester) {
        var suggestionName = "suggest";
        var query = new NativeQuery(new NativeQueryBuilder()
                .withSuggester(Suggester.of(builder -> builder.suggesters(suggestionName, suggester))));
        var suggest = operations.search(query, entityClass)
                .getSuggest();
        return Objects.isNull(suggest) ? Collections.emptyList() : suggest
                .getSuggestion(suggestionName)
                .getEntries()
                .stream()
                .map(Suggest.Suggestion.Entry::getText)
                .toList();
    }

    @Override
    public QueryPage<T> queryForPage(Query searchQuery) {
        var hits = operations.search(searchQuery, entityClass,
                entityInformation.getIndexCoordinates());
        var pageable = searchQuery.getPageable();
        return new SimpleQueryPage<>(hits.get()
                .map(SearchHit::getContent)
                .toList(), hits.getTotalHits(),
                pageable.getPageNumber(), pageable.getPageSize());
    }
}
